# encoding: utf-8

from typing import List
from pydantic import BaseModel, Field
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from train_with_torch import Model

import uvicorn
from fastapi import FastAPI

app = FastAPI()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using device:", device)
weights_path = "./weights/model_weights.pth"
model = Model()
model.load_state_dict(torch.load(weights_path, map_location=torch.device(device)))

model.to(device)
model.eval()


class FormModel(BaseModel):
    days: int = Field(default=1, description="日期")
    history: List[float] = Field(...)


class MyDataset(Dataset):
    def __init__(self, data_list: List[float]):
        self.data = np.array([data_list])

    def __getitem__(self, item):
        return self.data[item]

    def __len__(self):
        return len(self.data)


def predict_next_day(history: List[float]) -> float:
    if len(history) > 10:
        history = history[-10:]
    data_loader = DataLoader(MyDataset(history), batch_size=1, shuffle=False)
    with torch.no_grad():
        for data in data_loader:
            data = data.float()
            data = data.to(device)
            output = model(data)
            output_float = output.cpu().data.item()
            print("output_float:", output_float)
            return output_float


def predict_stock_price(history: List[float], days: int = 1) -> List[float]:

    result = []

    for day in range(days):
        next_day = predict_next_day(history)
        next_day = round(next_day, 2)
        result.append(next_day)
        history.append(next_day)

    return result


@app.post("/predict/stock-price")
def predict(form: FormModel):
    """
    预测股票价格
    """
    result = predict_stock_price(form.history, form.days)
    return result


if __name__ == '__main__':
    uvicorn.run(app, port=10054, host="0.0.0.0")
